#!/usr/bin/env python3
# _*_ coding:utf-8 _*_
'''
4. 使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。(2)
   (优化角点更新)
'''
import cv2 as cv
import numpy as np
# 加载视频
# Videofilename = r'D:\OpenCV-4.2.0-vc14_vc15\opencv\sources\samples\data\vtest.avi'
cap = cv.VideoCapture('vtest.avi')
if not cap.isOpened():
    print("无法打开视频文件")

# 角点检测参数设定
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)

# lucas_kanade光流法参数设定
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))

track_len = 10
detect_interval = 5
tracks = []
frame_idx = 0

flag, source = cap.read()                                   # 读取视频帧
gray = cv.cvtColor(source, cv.COLOR_BGR2GRAY)               # 转化为灰度虚图像
while True:
    ret, frame = cap.read()                                 # 读取视频帧
    if ret:
        frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)  # 转化为灰度虚图像
        vis = frame.copy()

        if len(tracks) > 0:                                 # 检测到角点后进行光流跟踪
            p0 = np.float32([tr[-1] for tr in tracks]).reshape(-1, 1, 2)
            # 前一帧的角点和当前帧的图像作为输入来得到角点在当前帧的位置
            p1, st, err = cv.calcOpticalFlowPyrLK(gray, frame_gray, p0, None,**lk_params)
            # 当前帧跟踪到的角点及图像和前一帧的图像作为输入来找到前一帧的角点位置
            p0r, st, err = cv.calcOpticalFlowPyrLK(frame_gray, gray, p1, None,**lk_params)

            d = abs(p0 - p0r).reshape(-1, 2).max(-1)                            # 得到角点回溯与前一帧实际角点的位置变化关系
            good = d < 1                                                        # 判断d内的值是否小于1，大于1跟踪被认为是错误的跟踪点
            new_tracks = []
            for tr, (x, y), good_flag in zip(tracks, p1.reshape(-1, 2), good):  # 将跟踪正确的点列入成功跟踪点
                if not good_flag:
                    continue
                tr.append((x, y))
                if len(tr) > track_len:
                    del tr[0]
                new_tracks.append(tr)
                cv.circle(vis, (x, y), 2, (0, 255, 0), -1)
            tracks = new_tracks
            cv.polylines(vis, [np.int32(tr) for tr in tracks], False,(0, 255, 0))  # 以上一振角点为初始点，当前帧跟踪到的点为终点划线

        if frame_idx % detect_interval == 0:                                     # 每5帧检测一次特征点
            mask = np.zeros_like(frame_gray)                                     # 初始化和视频大小相同的图像
            mask[:] = 255                                                        # 将mask赋值255也就是算全部图像的角点
            for x, y in [np.int32(tr[-1]) for tr in tracks]:                     # 跟踪的角点画圆
                cv.circle(mask, (x, y), 5, 0, -1)
            p = cv.goodFeaturesToTrack(frame_gray, mask=mask, **feature_params)  # 像素级别角点检测
            if p is not None:
                for x, y in np.float32(p).reshape(-1, 2):
                    tracks.append([(x, y)])  # 将检测到的角点放在待跟踪序列中

        frame_idx += 1
        gray = frame_gray.copy()
        cv.imshow('lk_track', vis)

    ch = 0xFF & cv.waitKey(100)
    if ch == 27:
        break
